package tjs.study.notes.dotest.jvm.classload.agent.impl;

import org.apache.commons.lang3.reflect.FieldUtils;
import org.apache.commons.lang3.reflect.MethodUtils;

import java.lang.reflect.InvocationTargetException;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.*;
import java.util.jar.JarFile;

public class ClassIsolationLoader extends URLClassLoader {

    protected List<ClassLoader> dependenceClassLoaders = new ArrayList<ClassLoader>();
    protected Set<String> includePackages = new HashSet<String>();
    protected Set<String> excludePackages = new HashSet<String>();

    protected Set<String> includeClasses = new HashSet<String>();
    protected Set<String> excludeClasses = new HashSet<String>();

    public ClassIsolationLoader(URL[] urls, ClassLoader parent) {
        super(urls, parent);
    }

    public static Builder builder(URL[] urls, ClassLoader parent) {
        return new ClassIsolationLoader(urls, parent).new Builder();
    }

    public static Builder builder(URL[] urls) {
        return builder(urls, null);
    }

    @Override
    protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {

        if (!isExclude(name) && isInclude(name)) {
            // 指定从父加载器加载
            return super.loadClass(name, resolve);
        }
        for (ClassLoader dependenceClassLoader : dependenceClassLoaders) {
            try { // 优先从依赖的加载器加载
                return dependenceClassLoader.loadClass(name);
            } catch (ClassNotFoundException e) {
                // ignore
            }
        }
        try {
            Class<?> c = findLoadedClass(name);
            if (c != null) {
                return c;
            }
            c = findClass(name);
            if (c != null) {
                return c;
            }
        } catch (ClassNotFoundException e) {
            // ignore
        }
        return super.loadClass(name, resolve);
    }

    protected boolean isInclude(String name) {
        if (includeClasses.contains(name)) {
            return true;
        }
        for (String includePackage : includePackages) {
            if (name.startsWith(includePackage)) {
                return true;
            }

        }
        return false;
    }

    protected boolean isExclude(String name) {
        if (excludeClasses.contains(name)) {
            return true;
        }
        for (String includePackage : excludePackages) {
            if (name.startsWith(includePackage)) {
                return true;
            }
        }
        return false;
    }

    public void release() {
        try {
            Set<String> closeJarFiles = new HashSet<String>();
            final Object ucp = FieldUtils.readField(this, "ucp", true);
            final List<Object> loaders = (List<Object>) FieldUtils.readField(ucp, "loaders", true);
            for (int i = 0; loaders != null && i < loaders.size(); i++) {
                Object jar;
                try {
                    jar = FieldUtils.readField(loaders.get(i), "jar", true);
                } catch (Exception e) {
                    continue;
                }
                if (jar instanceof JarFile) {
                    closeJarFiles.add(((JarFile) jar).getName());
                    ((JarFile) jar).close();
                }
            }
            if (!closeJarFiles.isEmpty()) {
                final Class<?> jarURLConnectionClass = Class.forName("sun.net.www.protocol.jar.JarURLConnection");
                final Object factory = FieldUtils.readStaticField(jarURLConnectionClass, "factory", true);
                final Map fileCache = (Map) FieldUtils.readField(factory, "fileCache", true);
                final Map urlCache = (Map) FieldUtils.readField(factory, "urlCache", true);
                Set<JarFile> jarFiles = new HashSet<JarFile>();
                if (fileCache != null) {
                    jarFiles.addAll(fileCache.values());
                }
                if (urlCache != null) {
                    jarFiles.addAll(urlCache.keySet());
                }
                Set<JarFile> removeJarFiles = new HashSet<JarFile>(closeJarFiles.size());
                for (JarFile jarFile : jarFiles) {
                    if (closeJarFiles.contains(jarFile.getName())) {
                        jarFile.close();
                        removeJarFiles.add(jarFile);
                    }
                }
                fileCache.values().removeAll(removeJarFiles);
                urlCache.keySet().removeAll(removeJarFiles);
            }
            final Vector nativeLibraries = (Vector) FieldUtils.readField(this, "nativeLibraries", true);
            if (nativeLibraries != null) {
                for (Object nativeLibrary : nativeLibraries) {
                    MethodUtils.invokeMethod(nativeLibrary, "finalize", null);
                }
            }
        } catch (Exception ex) {
            // ignore
        } finally {
            try {
                MethodUtils.invokeMethod(this, "close", null);
            } catch (Exception ex) {
                // ignore
            }
        }
    }

    public class Builder {
        public Builder dependenceClassLoaders(List<ClassLoader> dependenceClassLoaders) {
            ClassIsolationLoader.this.dependenceClassLoaders = dependenceClassLoaders;
            return this;
        }

        public Builder includePackages(Set<String> includePackages) {
            ClassIsolationLoader.this.includePackages = includePackages;
            return this;
        }

        public Builder excludePackages(Set<String> excludePackages) {
            ClassIsolationLoader.this.excludePackages = excludePackages;
            return this;
        }

        public Builder includeClasses(Set<String> includeClasses) {
            ClassIsolationLoader.this.includeClasses = includeClasses;
            return this;
        }

        public Builder excludeClasses(Set<String> excludeClasses) {
            ClassIsolationLoader.this.excludeClasses = excludeClasses;
            return this;
        }

        public ClassIsolationLoader build() {
            return ClassIsolationLoader.this;
        }
    }
}